package com.lcf.app.generator;

import com.baomidou.mybatisplus.generator.config.OutputFile;
import com.baomidou.mybatisplus.generator.config.po.TableField;
import com.baomidou.mybatisplus.generator.config.po.TableInfo;
import com.baomidou.mybatisplus.generator.engine.FreemarkerTemplateEngine;

import java.io.File;
import java.util.*;

/**
 * 代码生成器支持自定义[DTO\VO等]模版
 */
public class EnhanceFreemarkerTemplateEngine extends FreemarkerTemplateEngine {

    @Override
    protected void outputCustomFile( Map<String, String> customFile,  TableInfo tableInfo,  Map<String, Object> objectMap) {
        /**** 定义DTO，VO 配置  start****/
        TableInfo table = (TableInfo) objectMap.get("table");
        Set<String> dtoImportPackages = new HashSet<>();
        Iterator<String> iterator = table.getImportPackages().iterator();
        while (iterator.hasNext()) {
            String next = iterator.next();
            boolean b = next.indexOf("com.baomidou.mybatisplus.annotation") == -1;
            if (b) {
                dtoImportPackages.add(next);
            }
        }
        objectMap.put("dtoImportPackages", dtoImportPackages);
        //dto不需要update_time和create_time 字段
        List<TableField> dtoFields = new ArrayList<>();
        table.getFields().forEach(field ->{
            String name = field.getName();
            boolean b = !name.equals("update_time") && !name.equals("create_time");//dto不需要update_time和create_time 字段
            if (b) {
                dtoFields.add(field);
            }
        });
        //dto
        objectMap.put("dtoFields", dtoFields);

        /**** 定义DTO，VO 配置  end****/

        String entityName = tableInfo.getEntityName();
        String otherPath = this.getPathInfo(OutputFile.other);
        customFile.forEach((key, value) -> {
            String fileName = String.format(otherPath + File.separator + entityName + "%s", key);
            this.outputFile(new File(fileName), objectMap, value);
        });
        //
    }
}